import numpy as np
import iio

import time
from services.ad936x.api import Ad936xApi, DoubleGain, DoubleFrequency
from services.server_integration.api import ServerIntegrationApi
from common.third_party.devmem import DevMem, DevMemCache


AD_API = Ad936xApi()
SI_API = ServerIntegrationApi()


def sync_ad3(index: int):
    SERVICE_CLOCK = 0x43C5_0000
    mem = DevMemCache()

    ad_index = [0, 0, 1, 1, 2, 2][index]
    gain = [10, 10, 10, 10, 10, 10][index]
    config = AD_API.get_config(ad_index)
    AD_API.set_gain(3, DoubleGain(gain_0=gain, gain_1=gain))
    AD_API.set_center_frequency(3, DoubleFrequency(
        rx_frequency=config.tx_frequency,
        tx_frequency=config.tx_frequency
    ))

    multiplexer_state = [4, 5, 3, 2, 1, 0][index]
    mem.get(SERVICE_CLOCK, 1)\
        .write(0, multiplexer_state)


def services_pilot_set(index: int, state: bool):
    bin_index = [1, 3, 4, 8][index]
    mem = DevMem(address=0x83C1_0000, word_size=2)
    mem[bin_index] = state

def find_device():
    ctx = iio.Context()
    for dev in ctx.devices:
        if len(dev.buffer_attrs) > 0:
            return dev

    return None


dev: iio.Device = find_device()
dev.set_kernel_buffers_count(1)
for ch in dev.channels:
    ch.enabled = True

buff = iio.Buffer(dev, 2048*2)


def read_waveform(index: int):
    DevMem(0x83C2_0000, word_size=2)\
        .write(0, index)

    for _ in range(1):
        DevMem(0x83C1_0000, word_size=4)[2] = True
        DevMem(0x83C1_0000, word_size=4)[2] = False

        time.sleep(0.000_24)

        buff.refill()

    data = np.frombuffer(buff.read(), dtype=np.uint16)[:2048]
    return data[0::2], data[1::2]


with open('SERVICE_0_TX.txt', 'wt') as file:
    read_waveform(4)
    file.write('============================================\n')
    for _ in range(16):
        i_data, q_data = read_waveform(4)
        for i, q in zip(i_data, q_data):
            file.write(f'{i}\t{q}\n')

        file.write('============================================\n')


